# -*- coding: utf-8 -*-
'''
@author : 景色
@csdn : https://blog.csdn.net/pjjing
@QQ群 : 767101469
@公众号 : QuantRoad2019
'''
import thostmduserapi as mdapi
import csv, pandas as pd, os
from datetime import time

# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
'''
以下为需要订阅行情的合约号，注意选择有效合约；
有效连上但没有行情可能是过期合约或者不再交易时间内导致
'''
# subID=["au2104","IC2012","i2103","TA109"]

fut_list = pd.read_csv("../fut_list/all_futures.csv")
fut_list = fut_list[(fut_list.end_date >= "20201102") & (fut_list.start_date <= "20201102")]
fut_list = fut_list['name'].values.tolist()
subID = fut_list

class CBarData:
    def _init__(self):
        self.instrumentID: str
        self.updateTime: time
        self.volume: float = 0
        self.openInterest: float = 0
        self.openPrice: float = 0
        self.highPrice: float = 0
        self.lowPrice: float = 0
        self.closePrice: float = 0

class CFtdcMdSpi(mdapi.CThostFtdcMdSpi):

    def __init__(self, tapi):
        mdapi.CThostFtdcMdSpi.__init__(self)
        self.tapi = tapi
        self.csvfiles = {}
        self.submap = {}
        self.bar = None
        self.lastVolume: double = 0
    def __openCsv(self):
        csvheader = (["date", "open", "high", "low", "close"])
        for id in subID:
            csvname = f"../tick_data/{id}.csv"
            csvfile = open(csvname, 'a+', newline='')
            csvfile_w = csv.writer(csvfile)
            csvfile_w.writerow(csvheader)
            self.csvfiles[id] = csvfile
            self.submap[id] = csvfile_w

    def OnFrontConnected(self) -> "void":
        print("OnFrontConnected")
        loginfield = mdapi.CThostFtdcReqUserLoginField()
        # loginfield.BrokerID="1080"
        # loginfield.UserID="901206991"
        # loginfield.Password="smzy1108"
        loginfield.BrokerID = "9999"
        loginfield.UserID = "171525"
        loginfield.Password = "1234qwer@"
        loginfield.UserProductInfo = "python dll"
        self.tapi.ReqUserLogin(loginfield, 0)
        self.__openCsv()

    def OnRspUserLogin(self, pRspUserLogin: 'CThostFtdcRspUserLoginField', pRspInfo: 'CThostFtdcRspInfoField',
                       nRequestID: 'int', bIsLast: 'bool') -> "void":
        print(
            f"OnRspUserLogin, SessionID={pRspUserLogin.SessionID},ErrorID={pRspInfo.ErrorID},ErrorMsg={pRspInfo.ErrorMsg}")
        ret = self.tapi.SubscribeMarketData([id.encode('utf-8') for id in subID], len(subID))

    def OnRtnDepthMarketData(self, pDepthMarketData: 'CThostFtdcDepthMarketDataField') -> "void":
        print("OnRtnDepthMarketData")
        # 根据行情中的UpdateTime字段判断是否为新1分钟
        st = pDepthMarketData.UpdateTime.split(':')
        if not self.bar:
            newMinitue = True
        else:
            if int(st[1]) == self.bar.updateTime.minute:
                newMinitue = False
            else:
                newMinitue = True
        # 如果是新1分钟，生成一个新k线变量，CBarData结构体中有OHLC,time等K线字段
        if newMinitue:
            self.bar = CBarData()
            self.bar.instrumentID = pDepthMarketData.InstrumentID
            self.bar.exchangeID = pDepthMarketData.ExchangeID
            self.bar.updateTime = time(int(st[0]), int(st[1]), 0, 0)
            self.bar.volume = 0
            self.bar.openInterest = pDepthMarketData.OpenInterest
            self.bar.openPrice = pDepthMarketData.LastPrice
            self.bar.highPrice = pDepthMarketData.LastPrice
            self.bar.lowPrice = pDepthMarketData.LastPrice
            self.bar.closePrice = pDepthMarketData.LastPrice
            mdlist = ([pDepthMarketData.UpdateTime], [pDepthMarketData.LastPrice], [pDepthMarketData.LastPrice], [pDepthMarketData.LastPrice], [pDepthMarketData.LastPrice])
        else:
            # 如果不是新1分钟，将最新价与HL价相比然后更新，跟新C价
            self.bar.highPrice = max(self.bar.highPrice, pDepthMarketData.LastPrice)
            self.bar.lowPrice = min(self.bar.lowPrice, pDepthMarketData.LastPrice)
            mdlist = ([pDepthMarketData.UpdateTime], [pDepthMarketData.LastPrice], [self.bar.highPrice],
                      [self.bar.lowPrice], [pDepthMarketData.LastPrice])
        print(pDepthMarketData.InstrumentID)
        print(mdlist)
        self.submap[pDepthMarketData.InstrumentID].writerow(mdlist)
        self.csvfiles[pDepthMarketData.InstrumentID].flush()

    def OnRspSubMarketData(self, pSpecificInstrument: 'CThostFtdcSpecificInstrumentField',
                           pRspInfo: 'CThostFtdcRspInfoField', nRequestID: 'int', bIsLast: 'bool') -> "void":
        print("OnRspSubMarketData")
        print("InstrumentID=", pSpecificInstrument.InstrumentID)
        print("ErrorID=", pRspInfo.ErrorID)
        print("ErrorMsg=", pRspInfo.ErrorMsg)


def main():
    mduserapi = mdapi.CThostFtdcMdApi_CreateFtdcMdApi()
    mduserspi = CFtdcMdSpi(mduserapi)
    '''以下是7*24小时环境'''
    # mduserapi.RegisterFront("tcp://180.169.112.52:42213")
    mduserapi.RegisterFront("tcp://218.202.237.33:10112")
    mduserapi.RegisterSpi(mduserspi)
    mduserapi.Init()
    mduserapi.Join()


if __name__ == '__main__':
    main()